#include <openssl/evp.h>
#include <openssl/err.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

void handleErrors(void) {
    fprintf(stderr, "An error occurred.\n");
    ERR_print_errors_fp(stderr);
    exit(1);
}

int read_key_and_iv(const char *key_file, const char *iv_file, unsigned char *key, unsigned char *iv) {
    FILE *kf = fopen(key_file, "rb");
    FILE *ivf = fopen(iv_file, "rb");
    if (!kf || !ivf) {
        fprintf(stderr, "Could not open key or IV file.\n");
        return -1;
    }
    size_t key_len = fread(key, 1, 16, kf);
    if (key_len != 16) {
        fprintf(stderr, "Key must be 16 bytes long.\n");
        fclose(kf);
        fclose(ivf);
        return -1;
    }
    size_t iv_len = fread(iv, 1, 16, ivf);
    if (iv_len != 16) {
        fprintf(stderr, "IV must be 16 bytes long.\n");
        fclose(kf);
        fclose(ivf);
        return -1;
    }
    fclose(kf);
    fclose(ivf);
    return 0;
}

int main(int argc, char *argv[]) {
    if (argc != 5) {
        fprintf(stderr, "Usage: %s <key_file> <iv_file> <input_file> <output_file>\n", argv[0]);
        return 1;
    }

    unsigned char key[16], iv[16];
    if (read_key_and_iv(argv[1], argv[2], key, iv) != 0) {
        return 1;
    }

    FILE *f_input = fopen(argv[3], "rb");
    FILE *f_output = fopen(argv[4], "wb");
    if (!f_input || !f_output) {
        fprintf(stderr, "Could not open input or output file.\n");
        return 1;
    }

    EVP_CIPHER_CTX *ctx = EVP_CIPHER_CTX_new();
    if (!ctx) handleErrors();
    
    if (1 != EVP_EncryptInit_ex(ctx, EVP_sm4_cbc(), NULL, key, iv))
        handleErrors();

    unsigned char buffer[1024], ciphertext[1024 + EVP_MAX_BLOCK_LENGTH];
    int bytes_read, ciphertext_len, final_len;

    while ((bytes_read = fread(buffer, 1, sizeof(buffer), f_input)) > 0) {
        if (1 != EVP_EncryptUpdate(ctx, ciphertext, &ciphertext_len, buffer, bytes_read))
            handleErrors();
        fwrite(ciphertext, 1, ciphertext_len, f_output);
    }

    if (1 != EVP_EncryptFinal_ex(ctx, ciphertext + ciphertext_len, &final_len))
        handleErrors();
    fwrite(ciphertext, 1, ciphertext_len + final_len, f_output);

    EVP_CIPHER_CTX_free(ctx);
    fclose(f_input);
    fclose(f_output);
    printf("Encryption complete.\n");
    return 0;
}
